## This module contains routines to handle genetic trees:
## - generate trees (random expressions in list form)
## - evaluate trees
## - tree manipulation as inserting a tree in another
## - support routines for internal use

import random
import types
import copy
from numpy import linspace

## Generates a random expression (tree) with given function set (fs), "nodes",
## and terminal set (ts), "leaves", with either method "full" or "grow". The
## tree depth is given by maxd
def generateRandomExpr(fs, ts, maxd=5, method='full'):
    if (maxd == 0 or (method == 'grow' and random.random() < 0.5)):
        expr = random.choice(ts)
    else:
        func = random.choice(fs)
        arg1 = generateRandomExpr(fs, ts, maxd-1, method)
        arg2 = generateRandomExpr(fs, ts, maxd-1, method)
        expr = [func,arg1,arg2]
    return expr

def evalTree(tree, x=0):
    if type(tree) == types.ListType:
        if tree[0] == '+':
            return evalTree(tree[1], x) + evalTree(tree[2], x)
        elif tree[0] == '-':
            return evalTree(tree[1], x) - evalTree(tree[2], x)
        elif tree[0] == '*':
            return evalTree(tree[1], x) * evalTree(tree[2], x)
        elif tree[0] == '%':
            arg2 = evalTree(tree[2], x)
            if arg2 != 0:
                return evalTree(tree[1], x) / arg2
            return 1
    elif type(tree) == types.StringType:
        ##return ds.getVar(tree)
        return x
    elif type(tree) == types.FloatType or type(tree) == types.IntType:
        return tree
    else:
        print "unknown tree element"

def calcDepth(tree, level=0):
    if type(tree) == types.ListType:
        d1 = calcDepth(tree[1],level+1)
        d2 = calcDepth(tree[2],level+1)
        return max(d1,d2)
    else:
        return level


def getRandomSubtree(tree):
    # note 1
    #  the function makes sure to give always a <list> as return value
    #  this means to stop earlier in some cases when the chosen branch
    #  has little depth than maximum tree depth
    # note 2
    #  the returned list is a (deep) copy of the origin
    maxdepth = random.randint(1,calcDepth(tree))

    for i in range(maxdepth):
        if type(tree) == types.ListType:
            if type(tree[1]) ==  types.ListType and type(tree[2]) == types.ListType:
                tree = tree[random.randint(1,2)]
            elif type(tree[1]) == types.ListType:
                tree = tree[1]
            elif type(tree[2]) == types.ListType:
                tree = tree[2]

    return copy.deepcopy(tree)


def insertSubtreeInTree(tree, subtree):
    def modifySubtree(root, subtree):
        choice = random.randint(0,2)
        if choice == 0:
            for i in range(3):
                root[i] = subtree[i]
        else:
            root[choice] = subtree

    def gotoNextSubtree(root, subtree, depth):
        if depth>0:
            if type(root[1]) == types.ListType and type(root[2]) == types.ListType:
                gotoNextSubtree(root[random.randint(1,2)],subtree,depth-1)
            elif type(root[1]) == types.ListType:
                gotoNextSubtree(root[1],subtree,depth-1)
            elif type(root[2]) == types.ListType:
                gotoNextSubtree(root[2],subtree,depth-1)
            else:
                # in case of tree end reached before maximum allowed depth;
                # whole subtree or part of subtree is replaced
                modifySubtree(root, subtree)
        else:
            # reached maximum depth within tree; def: whole subtree or part
            # of subtree is replaced
            modifySubtree(tree, subtree)

    maxdepth = random.randint(1,calcDepth(tree))
    gotoNextSubtree(tree,subtree,maxdepth)



if __name__ == '__main__':
    print "genotree test procedure"
    fs = ['+','-','*','%']
    ts = [1.0,2.0,3.0,4.0,'x']

    for i in range(20):
        gp = generateRandomExpr(fs, ts, 3, 'grow')
        print gp
        Y = []
        for x in X:
            Y.append(evalTree(gp, x))

    

    rplot = roby.r['plot']

    rplot(X,Y)
    


##        population.append(gp)
##        labels.append(lattice.parseTreeToLattice(gp))

##    import tkrenderer
##    tkrenderer.drawPopulation(100,population,10)
##    drawTree(lattice.parseTreeToLattice(gp))

##    def x():
##        return 1.3456

##    fs = ['+','-','*','%']
##    ts = [-1, 1, 2.5, 5, x]
##
##    from dataset import DataSet
##    ds = DataSet('simple.xls', ['x','y','z'])
##
##    tr1 = ['+',-1,['-',1,['*',2.5,['%',5,'x']]]]
##
##    x = []
##    fx = []
##    while 1:
##        #print "x=",ds.getVar('x'),"   f(x)=", evalTree(tr1, ds)
##        x.append(ds.getVar('x'))
##        fx.append(evalTree(tr1, ds))
##        if not ds.nextItem():
##            break

##    import matplotlib.pyplot as plt
##    plt.plot(x,fx)
##    plt.show()




